# @FileName  : main.py
# @Time      : 2025/2/25 11:18
# @Author    : LuZhaoHui
# @Software  : PyCharm
import os.path

import torch
from matplotlib import pyplot as plt
import cv2

def testGpu():
    # CUDA版本
    print(torch.version.cuda)
    # cuDNN版本
    print(torch.backends.cudnn.version())
# '''
import paddle
import numpy as np
from paddle.vision.transforms import Normalize


def testModel():
    transform = Normalize(mean=[127.5], std=[127.5], data_format="CHW")
    # 下载数据集并初始化 DataSet
    train_dataset = paddle.vision.datasets.MNIST(mode="train", transform=transform)
    test_dataset = paddle.vision.datasets.MNIST(mode="test", transform=transform)

    # 打印数据集里图片数量
    print(
        "{} images in train_dataset, {} images in test_dataset".format(
            len(train_dataset), len(test_dataset)
        )
    )
    # 模型组网并初始化网络
    lenet = paddle.vision.models.LeNet(num_classes=10)
    model = paddle.Model(lenet)

    # 模型训练的配置准备，准备损失函数，优化器和评价指标
    model.prepare(
        paddle.optimizer.Adam(parameters=model.parameters()),
        paddle.nn.CrossEntropyLoss(),
        paddle.metric.Accuracy(),
    )

    # 模型训练
    model.fit(train_dataset, epochs=5, batch_size=64, verbose=1)
    # 模型评估
    model.evaluate(test_dataset, batch_size=64, verbose=1)

    # 保存模型
    # model.save("./output/mnist")
    # 加载模型
    model.load("output/mnist")

    # 从测试集中取出一张图片
    img, label = test_dataset[1]
    # 将图片shape从1*28*28变为1*1*28*28，增加一个batch维度，以匹配模型输入格式要求
    img_batch = np.expand_dims(img.astype("float32"), axis=0)

    # 执行推理并打印结果，此处predict_batch返回的是一个list，取出其中数据获得预测结果
    out = model.predict_batch(img_batch)[0]
    pred_label = out.argmax()
    print("true label: {}, pred label: {}".format(label[0], pred_label))
    # 可视化图片
    
    plt.imshow(img[0])
# '''

def showImage(imageName):
    if os.path.exists(imageName):
        img = cv2.imread(imageName)
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        plt.imshow(img)
        plt.show()
    else:
        print("{} not exist".format(imageName))

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    testGpu()
    showImage('e:/Image/ocrTest/id0.jpg')

    # print('paddleOcr test ver=%s' % (paddle.__version__))
    # testModel()

# See PyCharm help at https://www.jetbrains.com/help/pycharm/
